import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report
from joblib import dump

# 读取数据文件
df = pd.read_csv('daily.csv')

# 计算收益与风险比率

df['max_ratio'] = df['max_close'] / df['the_close'] # 收益潜力
df['min_ratio'] = df['min_close'] / df['the_close'] # 风险程度

# 自动划分高低阈值
high_retuen_threshold = df['max_ratio'].quantile(0.4) # 前40%定义为高收益
high_risk_threshold = df['min_ratio'].quantile(0.4) # 前40%定义为高风险

# 生成分类标签
conditions = [
    (df['max_ratio'] >= high_retuen_threshold) & (df['min_ratio'] <= high_risk_threshold),
    (df['max_ratio'] >= high_retuen_threshold) & (df['min_ratio'] > high_risk_threshold),
    (df['max_ratio'] < high_retuen_threshold) & (df['min_ratio'] > high_risk_threshold),
    (df['max_ratio'] < high_retuen_threshold) & (df['min_ratio'] <= high_risk_threshold),
]

labels = ['高收益高风险', '高收益低风险', '低收益低风险', '低收益高风险']
df['category'] = np.select(conditions, labels, default='未知')

# 特征选择
features = df[[
    'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 
    'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta', 'roic'
]]

# 标签编码
from sklearn.preprocessing import LabelEncoder
le = LabelEncoder()
df['category_encoded'] = le.fit_transform(df['category'])

# 数据标准化
scaler = StandardScaler()
X = scaler.fit_transform(features)
y = df['category_encoded']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 训练KNN模型
knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(X_train, y_train)

# 预测与评估
y_pred = knn.predict(X_test)
print(classification_report(y_test, y_pred, target_names=le.classes_))

# 保存模型和预处理对象
dump(knn, 'knn_classifier.joblib')
dump(scaler, 'feature_scaler.joblib')
dump(le, 'label_encoder.joblie')
